#include "cuda_tracer.cuh"
#include "math/cuMath.cuh"


//////////////////////////////////////////////
/// Almost device part

#define	RECURSION_DEPTH		5
#define	MAX_RAYS			(1 << RECURSION_DEPTH)
#define	USE_ANTIALIASING	0

struct CudaScene
{
	__device__ CudaScene()
		: spheres(0)
		, numSpheres(0)
		, boxes(0)
		, numBoxes(0)
		, planes(0)
		, numPlanes(0)
		, materials(0)
		, lights(0)
		, numLights(0)
	{
	}

	const CudaSphere*		spheres;
	int						numSpheres;
	const CudaBox*			boxes;
	int						numBoxes;
	const CudaPlane*		planes;
	int						numPlanes;
	const CudaMaterial*		materials;
	const CudaLight*		lights;
	int						numLights;
};

#define	MAX_CUTRACER_TEXTURES	4
#define	DECLARE_TEXTURE(number)	\
	texture<uchar4, cudaTextureType2D, cudaReadModeNormalizedFloat>	g_Texture_##number;

DECLARE_TEXTURE(0)
DECLARE_TEXTURE(1)
DECLARE_TEXTURE(2)
DECLARE_TEXTURE(3)

__device__ int FindIntersection(const Ray& ray, const CudaScene& scene, RayHitInfo& hit, float maxDist)
{
	if (ray.weight < 0.01f)
		return 0;

	hit.distToHit = maxDist;
	int hitSphereId = -1;
	for (int i = 0; i < scene.numSpheres; ++i)
	{
		RayHitInfo hitInfo;
		if (SphereRayCollisionCheck(scene.spheres[i], ray, hitInfo, hit.distToHit))
		{
			hit = hitInfo;
			hitSphereId = i;
		}
	}

	int hitBoxId = -1;
	for (int i = 0; i < scene.numBoxes; ++i)
	{
		RayHitInfo hitInfo;
		if (BoxRayCollisionCheck(scene.boxes[i], ray, hitInfo, hit.distToHit))
		{
			hit = hitInfo;
			hitBoxId = i;
			hitSphereId = -1;
		}
	}

	int hitPlaneId = -1;
	for (int i = 0; i < scene.numPlanes; ++i)
	{
		RayHitInfo hitInfo;
		if (PlaneRayCollisionCheck(scene.planes[i], ray, hitInfo, hit.distToHit))
		{
			hit = hitInfo;
			hitPlaneId = i;
			hitBoxId = -1;
			hitSphereId = -1;
		}
	}

	if (hitSphereId != -1)
	{
		hit.matID = scene.spheres[hitSphereId].matID;
		return 1;
	}
	else if (hitBoxId != -1)
	{
		hit.matID = scene.boxes[hitBoxId].matID;
		return 1;
	}
	else if (hitPlaneId != -1)
	{
		hit.matID = scene.planes[hitPlaneId].matID;
		return 1;
	}
	else
		return 0;
}

__device__ __inline__ float4 SampleTexture(int texId, float s, float t)
{
	switch (texId)
	{
		case 0:
			return tex2D(g_Texture_0, s, t);
		case 1:
			return tex2D(g_Texture_1, s, t);
		case 2:
			return tex2D(g_Texture_2, s, t);
		case 3:
			return tex2D(g_Texture_3, s, t);
	}

	return make_float4(0.0f, 0.0f, 0.0f, 1.0f);
}

__device__ float4 Shader(const Ray& ray, const RayHitInfo& hitInfo, const CudaScene& scene)
{
	float4 color = make_float4(0.0f, 0.0f, 0.0f, 1.0f);

	const CudaMaterial& mat = scene.materials[hitInfo.matID];
	float colorInfluence = clamp(1.0f - mat.reflectivity - mat.transparency, 0.0f, 1.0f);

	float3 matDiffuse = mat.diffuse;
	if (mat.albedo != -1)
	{
		float4 texel = SampleTexture(mat.albedo, hitInfo.hitTexcoord.x, hitInfo.hitTexcoord.y);
		matDiffuse.x = texel.x;
		matDiffuse.y = texel.y;
		matDiffuse.z = texel.z;
	}

	for (int i = 0; i < scene.numLights; ++i)
	{
		const CudaLight& light = scene.lights[i];

		float3 dirToLight = float3_sub(light.position, hitInfo.hitPos);
		float distToLight = length(dirToLight);
		dirToLight = float3_div(dirToLight, distToLight);
		float NdotL = dot(hitInfo.hitNormal, dirToLight);

		Ray lightRay(hitInfo.hitPos, dirToLight);
		RayHitInfo lightHitInfo;
		if (!FindIntersection(lightRay, scene, lightHitInfo, distToLight))
		{
			float lambert = clamp(NdotL, 0.0f, 1.0f);
			color.x += lambert * light.intensity.x * matDiffuse.x * colorInfluence;
			color.y += lambert * light.intensity.y * matDiffuse.y * colorInfluence;
			color.z += lambert * light.intensity.z * matDiffuse.z * colorInfluence;

			// point light source: sample once for specular highlight
			if (mat.specularLevel > 0.0f)
			{
				float3 V = ray.dir;
				float3 R = normalize(reflect(dirToLight, hitInfo.hitNormal));
				float VdotR = dot(V, R);
				if (VdotR > 0.0f)
				{
					float spec = powf(VdotR, mat.glossiness) * lambert * mat.specularLevel;
					// add specular component to ray color
					color.x += spec * light.intensity.x * mat.specular.x;// * colorInfluence;
					color.y += spec * light.intensity.y * mat.specular.y;// * colorInfluence;
					color.z += spec * light.intensity.z * mat.specular.z;// * colorInfluence;
				}
			}
		}
	}

	return float4_mul(color, ray.weight);
}

__device__ float4 TraceRay(const Ray& ray, const CudaScene& scene)
{
	float4 color = make_float4(0.0f, 0.0f, 0.0f, 1.0f);

	Ray currRays[MAX_RAYS];
	Ray newRays[MAX_RAYS];

	currRays[0] = ray;

	int numCurrRays = 1;
	int level = 0;
	do {
		int numNewRays = 0;

		for (int currRay = 0; currRay < numCurrRays; ++currRay)
		{
			const Ray& tracedRay = currRays[currRay];

			RayHitInfo hitInfo;
			if (FindIntersection(tracedRay, scene, hitInfo, (float)INT_MAX))
			{
				const CudaMaterial& mat = scene.materials[hitInfo.matID];
				// if we have bump texture - then apply transformation on our normal
				if (mat.bump != -1)
				{
					float4 texel = SampleTexture(mat.bump, hitInfo.hitTexcoord.x, hitInfo.hitTexcoord.y);
					float3 n = normalize(float3_sub(float3_mul(make_float3(texel.x, texel.y, texel.z), 2.0f), 1.0f));
					hitInfo.hitNormal = transform_by_normal(n, hitInfo.hitNormal);
				}

				color = float4_add(color, Shader(tracedRay, hitInfo, scene));

				if (mat.reflectivity > CU_EPSILON)
				{
					newRays[numNewRays].SetDir(normalize(reflect(tracedRay.dir, hitInfo.hitNormal)));
					newRays[numNewRays].origin = float3_add(hitInfo.hitPos, float3_mul(tracedRay.dir, CU_EPSILON));
					newRays[numNewRays].weight = tracedRay.weight * mat.reflectivity;
					numNewRays++;
				}

				if (mat.transparency > CU_EPSILON)
				{
					float refrIndex;
					const float NdorRD = dot(hitInfo.hitNormal, tracedRay.dir);
					if (NdorRD > 0.0f)
					{
						hitInfo.hitNormal = float3_mul(hitInfo.hitNormal, -1.0f);
						refrIndex = mat.refractionIndex;
					}
					else
						refrIndex = 1.0f / mat.refractionIndex;

					const float NdotI = -dot(hitInfo.hitNormal, tracedRay.dir);
					float cosT2 = 1.0f + refrIndex * refrIndex * (NdotI * NdotI - 1.0f);
					if (cosT2 > 0.0f)
					{
						newRays[numNewRays].SetDir(float3_add(float3_mul(tracedRay.dir, refrIndex), float3_mul(hitInfo.hitNormal, refrIndex * NdotI - sqrtf(cosT2))));
						newRays[numNewRays].origin = float3_add(hitInfo.hitPos, float3_mul(tracedRay.dir, CU_EPSILON));
						newRays[numNewRays].weight = tracedRay.weight * mat.transparency;
						numNewRays++;
					}
				}
			}
		}

		// move new into current
		for (int newRay = 0; newRay < numNewRays; ++newRay)
			currRays[newRay] = newRays[newRay];

		numCurrRays = numNewRays;

		level++;
	} while (level < RECURSION_DEPTH && numCurrRays > 0);

	color = saturate(color);
	return color;
}

__device__ __inline__ Ray CalcRay(const CudaCamera* camera, float x, float y, float width, float height)
{
	Ray ray;
	ray.origin = camera->position;
	ray.SetDir(float3_add(camera->direction,
			   float3_sub(float3_mul(camera->U, (x / (width - 1.0f)) * 2.0f - 1.0f),
			   float3_mul(camera->V, (y / (height - 1.0f)) * 2.0f - 1.0f))));

	return ray;
}

__global__ void kernel_PerformRaytracing(uint* screen, uint screenW, uint screenH, const CudaSphere* spheres, int numSpheres,
										 const CudaBox* boxes, int numBoxes, const CudaPlane* planes, int numPlanes,
										 const CudaMaterial* materials, const CudaCamera* camera, const CudaLight* lights, int numLights)
{
	const uint ux = __umul24(blockIdx.x, blockDim.x) + threadIdx.x;
	const uint uy = __umul24(blockIdx.y, blockDim.y) + threadIdx.y;

	const float width = (float)screenW;
	const float height = (float)screenH;

	if (ux < screenW && uy < screenH)
	{
		float x = (float)ux + 0.5f;
		float y = (float)uy + 0.5f;

		CudaScene scene;
		scene.spheres = spheres;
		scene.numSpheres = numSpheres;
		scene.boxes = boxes;
		scene.numBoxes = numBoxes;
		scene.planes = planes;
		scene.numPlanes = numPlanes;
		scene.materials = materials;
		scene.lights = lights;
		scene.numLights = numLights;

#if		!USE_ANTIALIASING	// no AA
		{
			Ray ray = CalcRay(camera, x, y, width, height);
			float4 color = TraceRay(ray, scene);
			screen[uy * screenW + ux] = RGBAFloatToUint(color);
		}
#else	// USE_ANTIALIASING		// AA version
		{
#define	TRACE_AA(num, offX, offY, coeff)									\
	Ray ray##num = CalcRay(camera, x + (offX), y + (offY), width, height);	\
	float4 color##num = float4_mul(TraceRay(ray##num, scene), coeff);

			TRACE_AA(0, +0.0f, +0.0f, 0.25f)	// center pixel
			TRACE_AA(1, -0.5f, -0.5f, 0.1875f)	// left-upper corner of pixel
			TRACE_AA(2, -0.5f, +0.5f, 0.1875f)	// left-lower corner of pixel
			TRACE_AA(3, +0.5f, +0.5f, 0.1875f)	// right-lower corner of pixel
			TRACE_AA(4, +0.5f, -0.5f, 0.1875f)	// right-upper corner of pixel

			float4 color = float4_add(float4_add(float4_add(float4_add(color0, color1), color2), color3), color4);
			screen[uy * screenW + ux] = RGBAFloatToUint(color);
		}
#endif	// USE_ANTIALIASING
	}
}

//////////////////////////////////////////////
/// Almost host part

__host__ uint uDivUp(uint a, uint b)
{
	return (a % b != 0) ? (a / b + 1) : (a / b);
}

__host__ void CudaRayTrace(unsigned int* imgData, unsigned int imgWidth, unsigned int imgHeight, const CudaSphere* spheres, int numSpheres,
						   const CudaBox* boxes, int numBoxes, const CudaPlane* planes, int numPlanes, const CudaMaterial* materials,
						   const CudaCamera* camera, const CudaLight* lights, int numLights)
{
	dim3 blockSize(16, 16);
	dim3 gridSize(uDivUp(imgWidth, blockSize.x), uDivUp(imgHeight, blockSize.y));

	kernel_PerformRaytracing<<<gridSize, blockSize>>>(imgData, imgWidth, imgHeight, spheres, numSpheres, boxes, numBoxes,
													  planes, numPlanes, materials, camera, lights, numLights);
}

cudaArray*	g_TexturesArrays[MAX_CUTRACER_TEXTURES];
__host__ void CudaInitTexture(const void* imageData, int width, int height, int texID)
{
	const int size = width * height * sizeof(unsigned int);

	// copy image data to array
	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(8, 8, 8, 8, cudaChannelFormatKindUnsigned);
	cudaSafeCall(cudaMallocArray(&g_TexturesArrays[texID], &channelDesc, width, height)); 
	cudaSafeCall(cudaMemcpyToArray(g_TexturesArrays[texID], 0, 0, imageData, size, cudaMemcpyHostToDevice));

#define	SETUP_TEXTURE(number)	\
	g_Texture_##number.addressMode[0] = cudaAddressModeWrap;	\
	g_Texture_##number.addressMode[1] = cudaAddressModeWrap;	\
	g_Texture_##number.filterMode = cudaFilterModeLinear;		\
	g_Texture_##number.normalized = true;						// access with normalized floats texture coordinates

#define	BIND_TEXTURE(number)	\
	cudaSafeCall(cudaBindTextureToArray(g_Texture_##number, g_TexturesArrays[number]));	// Bind the array to the texture

	switch (texID)
	{
		case 0:
			SETUP_TEXTURE(0)
			BIND_TEXTURE(0)
			break;
		case 1:
			SETUP_TEXTURE(1)
			BIND_TEXTURE(1)
			break;
		case 2:
			SETUP_TEXTURE(2)
			BIND_TEXTURE(2)
			break;
		case 3:
			SETUP_TEXTURE(3)
			BIND_TEXTURE(3)
			break;
	}
}

__host__ void CudaFreeTexture(int texID)
{
	cudaSafeCall(cudaFreeArray(g_TexturesArrays[texID]));
}

